import numpy as np
from numpy.ma import cos
import matplotlib.pyplot as plt
import math
import time
import matplotlib
from matplotlib.animation import FuncAnimation

matplotlib.use('TkAgg')
matplotlib.rcParams['font.family'] = 'SimHei'
matplotlib.rcParams['axes.unicode_minus'] = False


class GeneticAlgorithmOptimizer:
    def __init__(self):
        # 算法参数
        self.DNA_SIZE = 24  # 编码长度
        self.CROSS_RATE = 0.5  # 交叉率
        self.MUTA_RATE = 0.015  # 变异率
        self.Iterations = 200  # 迭代次数
        self.X_BOUND = [1, 2]  # X区间
        self.Y_BOUND = [1, 2]  # Y区间

        # 测试的种群规模
        self.POP_SIZES = [20, 50, 100, 200, 500]

        # 创建可视化图形
        self.fig = plt.figure(figsize=(16, 8))
        self.setup_visualization()

        # 开启交互模式
        plt.ion()

    def setup_visualization(self):
        """初始化可视化设置"""
        # 适应度收敛曲线
        self.ax_convergence = self.fig.add_subplot(121)
        self.ax_convergence.set_xlabel('迭代次数')
        self.ax_convergence.set_ylabel('最佳适应度')
        self.ax_convergence.set_title('适应度收敛曲线')
        self.ax_convergence.grid(True)

        # 运行时间比较
        self.ax_runtime = self.fig.add_subplot(122)
        self.ax_runtime.set_xlabel('种群规模')
        self.ax_runtime.set_ylabel('运行时间 (秒)')
        self.ax_runtime.set_title('运行时间比较')
        self.ax_runtime.grid(True)

        plt.tight_layout()

    def fitness_function(self, x, y):
        """适应度函数 - Rastrigin函数变种"""
        return 20 + x ** 2 + y ** 2 - 10 * (cos(2 * math.pi * x) + cos(2 * math.pi * y))

    def decode_genome(self, population):
        """解码DNA序列为实际参数值"""
        x_pop = population[:, 1::2]  # 奇数列表示X
        y_pop = population[:, ::2]  # 偶数列表示y

        # 将二进制转换为实数
        x = x_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * \
            (self.X_BOUND[1] - self.X_BOUND[0]) + self.X_BOUND[0]
        y = y_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * \
            (self.Y_BOUND[1] - self.Y_BOUND[0]) + self.Y_BOUND[0]

        return x, y

    def calculate_fitness(self, population):
        """计算种群适应度（转换为最小化问题）"""
        x, y = self.decode_genome(population)
        values = self.fitness_function(x, y)
        return -(values - np.max(values)) + 0.0001  # 转换为最小化问题

    def selection(self, population, fitness):
        """选择操作 - 轮盘赌选择"""
        idx = np.random.choice(
            np.arange(population.shape[0]),
            size=population.shape[0],
            replace=True,
            p=fitness / fitness.sum()
        )
        return population[idx]

    def crossover_and_mutation(self, population):
        """交叉和变异操作"""
        new_population = []
        for individual in population:
            # 创建副本避免修改原个体
            child = individual.copy()

            # 交叉操作
            if np.random.rand() < self.CROSS_RATE:
                mate_idx = np.random.randint(population.shape[0])
                crossover_point1 = np.random.randint(0, self.DNA_SIZE * 2 - 1)
                crossover_point2 = np.random.randint(crossover_point1, self.DNA_SIZE * 2)
                child[crossover_point1:crossover_point2] = \
                    population[mate_idx][crossover_point1:crossover_point2]

            # 变异操作
            self.mutate(child)

            new_population.append(child)

        return np.array(new_population)

    def mutate(self, individual):
        """变异操作"""
        if np.random.rand() < self.MUTA_RATE:
            mutate_point = np.random.randint(0, len(individual))
            individual[mutate_point] = individual[mutate_point] ^ 1

    def run_optimization_with_visualization(self, pop_size):
        """带可视化的遗传算法优化"""
        # 初始化种群
        population = np.random.randint(2, size=(pop_size, self.DNA_SIZE * 2))
        best_fitness_history = []

        # 清空当前图形
        self.ax_convergence.clear()
        self.ax_convergence.set_xlabel('迭代次数')
        self.ax_convergence.set_ylabel('最佳适应度')
        self.ax_convergence.set_title(f'种群规模 {pop_size} 的适应度收敛曲线')
        self.ax_convergence.grid(True)

        start_time = time.time()

        for iteration in range(self.Iterations):
            # 遗传操作
            population = self.crossover_and_mutation(population)
            fitness = self.calculate_fitness(population)

            # 解码当前种群
            x, y = self.decode_genome(population)
            current_fitness = self.fitness_function(x, y)

            # 记录最佳适应度
            best_idx = np.argmin(current_fitness)
            best_fitness_history.append(current_fitness[best_idx])

            # 更新适应度曲线
            if iteration % 5 == 0 or iteration == self.Iterations - 1:
                self.ax_convergence.plot(range(iteration + 1), best_fitness_history, 'b-')
                plt.draw()
                plt.pause(0.01)

            # 选择新一代
            population = self.selection(population, fitness)

        runtime = time.time() - start_time

        return {
            'pop_size': pop_size,
            'best_fitness': best_fitness_history[-1],
            'runtime': runtime,
            'fitness_history': best_fitness_history
        }

    def update_runtime_comparison(self, results):
        """更新运行时间比较图"""
        self.ax_runtime.clear()
        pop_sizes = [r['pop_size'] for r in results]
        runtimes = [r['runtime'] for r in results]

        bars = self.ax_runtime.bar(range(len(pop_sizes)), runtimes,
                                   color=plt.cm.viridis(np.linspace(0, 1, len(pop_sizes))),
                                   width=0.6)

        # 添加数值标签
        for bar in bars:
            height = bar.get_height()
            self.ax_runtime.text(bar.get_x() + bar.get_width() / 2., height,
                                 f'{height:.2f}s',
                                 ha='center', va='bottom')

        self.ax_runtime.set_xticks(range(len(pop_sizes)))
        self.ax_runtime.set_xticklabels(pop_sizes)
        self.ax_runtime.set_xlabel('种群规模')
        self.ax_runtime.set_ylabel('运行时间 (秒)')
        self.ax_runtime.set_title('运行时间比较')
        self.ax_runtime.grid(True)

        plt.draw()

    def generate_analysis_report(self, results):
        """生成分析报告"""
        print("\n===== 种群规模影响分析 =====")
        print("1. 小规模种群(20-50):")
        print("   - 收敛速度快，但容易陷入局部最优")
        print("   - 计算时间短，适合简单问题或快速验证")

        print("\n2. 中等规模种群(100-200):")
        print("   - 平衡了收敛速度和解的质量")
        print("   - 在大多数问题上表现良好")

        print("\n3. 大规模种群(500):")
        print("   - 找到全局最优的概率最高")
        print("   - 但计算成本显著增加")
        print("   - 适合复杂多峰函数优化")

        # 找出最佳折中选择
        best_balance = min(results, key=lambda x: x['runtime'] * abs(x['best_fitness']))
        print(f"\n建议：根据测试结果，种群规模 {best_balance['pop_size']} "
              f"在运行时间({best_balance['runtime']:.2f}s)和解质量({best_balance['best_fitness']:.6f})"
              "之间取得了较好的平衡")

    def run_comparative_study(self):
        """运行比较研究"""
        results = []

        print("开始测试不同种群规模...")
        for pop_size in self.POP_SIZES:
            print(f"\n正在运行种群规模 {pop_size}...")
            result = self.run_optimization_with_visualization(pop_size)
            results.append(result)

            # 更新运行时间比较图
            self.update_runtime_comparison(results)

            print(f"结果: 最佳适应度={result['best_fitness']:.6f}, 运行时间={result['runtime']:.2f}秒")

        # 生成分析报告
        self.generate_analysis_report(results)

        # 关闭交互模式并保持图形显示
        plt.ioff()
        plt.show()


if __name__ == "__main__":
    optimizer = GeneticAlgorithmOptimizer()
    optimizer.run_comparative_study()